import numpy as np
import torch
from torch import nn

from model_wrapper import SimpleClassifyModelWrapper

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(3, 1)
        self.fc2 = nn.Linear(2, 1)

    def forward(self, x1, x2):
        x1 = self.fc1(x1)
        x2 = self.fc2(x2)
        return torch.cat((x1, x2), dim=1)
    


if __name__ == "__main__":
    X1 = np.random.randn(8, 3)
    X2 = np.random.randn(8, 2)
    y = np.random.randint(0, 2, 8)

    model = Model()
    wrapper = SimpleClassifyModelWrapper(model)
    wrapper.train((X1, X2), y)
